# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: xguo

import atexit
import json
import logging
import enum
from concurrent.futures import ThreadPoolExecutor

from tornado import gen

import coin.flow.topic_map as topic_map
from coin.base.timestamp import get_timestamp

from selenium import webdriver
from selenium.webdriver.common.desired_capabilities import DesiredCapabilities
from selenium.webdriver.chrome.options import Options

import chromedriver_binary  # Download chromedriver if not exist


def _prepare_driver(url):
  options = Options()
  options.add_argument('--headless')  # Disable chrome gui.
  options.add_argument('--disable-gpu')  # Last I checked this was necessary.
  options.add_argument('--mute-audio')  # Disable bitflyer trade sound.

  caps = DesiredCapabilities.CHROME
  caps['loggingPrefs'] = {'performance': 'ALL'}
  driver = webdriver.Chrome('chromedriver', chrome_options=options, desired_capabilities=caps)
  driver.implicitly_wait(10)
  driver.get(url)
  driver.implicitly_wait(10)
  return driver


class BitflyerChromeDriver(object):
  def __init__(self, message_handler):
    self._handler = message_handler
    self._driver = None
    self._executor = ThreadPoolExecutor(max_workers=1)
    self._stop = True
    self._url = 'https://lightning.bitflyer.jp/home/demo?lang=en'

  def handle_websocket_message(self, ioloop):
    log_entries = None
    try:
      log_entries = self._driver.get_log("performance")
    except OSError:
      # OSError: [Errno 99] Cannot assign requested address
      pass

    if log_entries is None:
      return

    for entry in log_entries:
      try:
        message_str = entry['message']
        message = json.loads(message_str)
        method = message['message']["method"]
        if method.lower() == "Network.webSocketFrameReceived".lower():
          payload_str = message['message']['params']['response']['payloadData']
          try:
            payload = json.loads(payload_str)
            ioloop.add_callback(self._handler, get_timestamp(), payload)
          except ValueError:
            continue
      except KeyError:
        pass

  def start(self, ioloop):
    self._stop = False
    self._driver = _prepare_driver(self._url)
    atexit.register(self.close_driver)
    while not self._stop:
      self.handle_websocket_message(ioloop)

  def stop(self):
    self._stop = True
    self.close_driver()
    self._executor.shutdown()

  def run(self, ioloop):
    self._executor.submit(self.start, ioloop)

  def close_driver(self):
    self._stop = True
    # noinspection PyBroadException
    try:
      self._driver.quit()
    except Exception:
      pass


class BitflyerSeleniumSubscriber(object):
  class State(enum.Enum):
    INIT = 1
    RUNNING = 2
    CLOSED = 3
    STOP = 4

  def __init__(self,
               queue,
               group_id: str,
               worker_id: str,
               products,
               close_callback=None,
               logger=None):
    self._logger = logger or logging.getLogger(__name__)
    self._close_callback = close_callback or (lambda: None)
    self._teardown_callback = []

    self._state = self.State.INIT
    self._ioloop = None

    self._publisher = None
    self._init_publisher(queue, group_id, worker_id)
    self._driver = BitflyerChromeDriver(self.on_message)

  def __del__(self):
    self.stop()

  def _add_teardown_callback(self, callback):
    self._teardown_callback.append(callback)

  def _init_publisher(self, queue, group_id: str, worker_id: str):
    topic_board_snapshot = topic_map.search_topic_id('bitflyer_ws_ticker_board_%s%s' %
                                                     (group_id, worker_id))[0]

    self._publisher = queue.get_writer(topic_id=topic_board_snapshot[1])
    self._add_teardown_callback(self._teardown_publisher)

  def _teardown_publisher(self):
    if self._publisher:
      self._publisher.close()
    self._publisher = None

  def start(self, ioloop):
    assert not self.stopped

    self._state = self.State.RUNNING
    self._ioloop = ioloop
    ioloop.add_callback(self._run_loop)

  def stop(self):
    if self._state == self.State.STOP:
      return

    if self._state == self.State.RUNNING:
      self._logger.info('Stopped')

    if self._driver:
      self._driver.stop()
    self._state = self.State.STOP
    self._tear_down()

  @property
  def stopped(self):
    return self._state == self.State.STOP

  def _tear_down(self):
    for callback in self._teardown_callback:
      # noinspection PyBroadException
      try:
        callback()
      except Exception:
        self._logger.exception('Teardown callback raised an exception')

  def _notify_close(self):
    if self._state != self.State.RUNNING:
      return

    self._state = self.State.CLOSED
    self._logger.error('Connection closed')

    # noinspection PyBroadException
    try:
      self._close_callback()
    except Exception:
      self._logger.exception('Close callback raised an exception')

  @gen.coroutine
  def _run_loop(self):
    if self._state != self.State.RUNNING:
      return

    self._driver.run(self._ioloop)
    try:
      while self._state == self.State.RUNNING:
        yield gen.sleep(1)

    finally:
      self._driver.stop()
      self._driver = None
      self._notify_close()

  def on_pubnub_error(self, pubnub_status):
    if self._state == self.State.RUNNING:
      self._logger.error('Disconnected. status_category=%d' % pubnub_status.category)
    self._state = self.State.CLOSED

  def on_message(self, timestamp: int, message):
    if self.stopped:
      return

    try:
      self._publisher.write(timestamp, message)
    except KeyError:
      self._logger.error('Unknown event_type: %s' % message)
